// Copyright (c) 1996-2002 John Lyon-Smith. All rights reserved.

#pragma once

#include <wxtl\Base.h>
#include <wxtl\CriticalSection.h>
#include <wxtl\RegistryKey.h>
#include <wxtl\ComClassFactory.h>
#include <wxtl\StringUtil.h>
#include <wxtl\Trace.h>
#include <unknwn.h>

using namespace boost;

// Here's a bunch of useful COM macros

#define THROW_BAD_HR(x)	{ register HRESULT hr = (x); if (FAILED(hr)) throw _com_error(hr); }

#define CLEANUP_BAD_HR(x, y)	{ register HRESULT hr = (x); if (FAILED(hr)) goto y; }

#define RET_BAD_HR(x)	{ register HRESULT hr = (x); if (FAILED(hr)) return hr; }
#define RET_GOOD_HR(x)	{ register HRESULT hr = (x); if (SUCCEEDED(hr)) return hr; }

#define ASSERT_GOOD_HR(x)	_ASSERT(SUCCEEDED(x))

#ifdef _DEBUG
#define VERIFY_GOOD_HR(x)	_ASSERT(SUCCEEDED(x))
#else // _DEBUG
#define VERIFY_GOOD_HR(x)	(x)
#endif // _DEBUG

#pragma warning(push)
#pragma warning(disable: 4239)	// non-standard extension

namespace wxtl
{
	inline tstring _stdcall stringFromGuid(REFGUID rguid)
	{
		TCHAR buf[512];
		
		_stprintf_s(buf,512, _T("{%08X-%04X-%04X-%02X%02X-%02X%02X%02X%02X%02X%02X}"), 
			rguid.Data1, rguid.Data2, rguid.Data3,
			(int)rguid.Data4[0], (int)rguid.Data4[1], (int)rguid.Data4[2], (int)rguid.Data4[3],
			(int)rguid.Data4[4], (int)rguid.Data4[5], (int)rguid.Data4[6], (int)rguid.Data4[7]);
			
		return tstring(buf);
	}

	// Load a library from a CLSID
	inline HRESULT _stdcall loadLibraryFromClsid(
		REFCLSID rclsid,
		HINSTANCE* phDll)
	{
		HRESULT hr = S_OK;

		// Convert the clsid to a string so we can find the InprocServer32 in the registry.
		tstring guidString(stringFromGuid(rclsid).c_str());
		tstring keyString(tstring(_T("CLSID\\")) + guidString + _T("\\InprocServer32"));
		tstring str;
		
		try
		{
			RegistryKey key(HKEY_CLASSES_ROOT);

			key = key.openKey(keyString.c_str());

			key.queryValue(NULL, str);
		}
		catch (_win32_error& e)
		{
			return HRESULT_FROM_WIN32(e.error());
		}	    
		
		// We've got the name of the DLL to load, so load it.
		*phDll = LoadLibraryEx(str.c_str(), NULL, LOAD_WITH_ALTERED_SEARCH_PATH);
		
		if (*phDll == NULL)
			hr = HRESULT_FROM_WIN32(GetLastError());
	    
		return hr;
	}

	// Use to create COM objects without invoking the wrath of COM
	inline HRESULT _stdcall fakeCoCreateInstance(
		REFCLSID rclsid,
		REFIID riid,
		void **ppv)
	{
		HINSTANCE hDll;
		HRESULT hr = loadLibraryFromClsid(rclsid, &hDll);

		if (FAILED(hr))
			return hr;
					
		// We've loaded the DLL, so find the DllGetClassObject function.
		FARPROC func = GetProcAddress(hDll, "DllGetClassObject");

		if (func != NULL)
		{
			// Cool, lets call the function to get a class factory for the rclsid passed in.
			IClassFactory *classFactory;
			PFN_DLLGETCLASSOBJECT *pfnDllGetClassObject = (PFN_DLLGETCLASSOBJECT*)func;

			hr = pfnDllGetClassObject(rclsid, IID_IClassFactory, (void**)&classFactory);

			if (SUCCEEDED(hr))
			{
				// Ask the class factory to create an instance of the necessary object.
				hr = classFactory->CreateInstance(NULL, riid, ppv);

				// Release that class factory!
				classFactory->Release();
			}
		}
		else
			hr = HRESULT_FROM_WIN32(GetLastError());

		return hr;
	}

	////////////////////////////////////////////////////////////////////////////////////////////////////

	class ComHelp;
	
	class ComHelp
	{
		//************************************
		// Method:    getComHelp
		// FullName:  wxtl::ComHelp::getComHelp
		// Access:    private 
		// Returns:   friend ComHelp& _stdcall
		// Qualifier: 
		//************************************
		//friend 

	public:
		ComHelp() {}
		~ComHelp() {}
		
	private:
		static bool registerComClass( 
			REFCLSID rclsid,
			PCTSTR szDesc,					
			PCTSTR szProgIDPrefix,
			PCTSTR szVersion,				
			PCTSTR szClassProgID,	
			PCTSTR szThreadingModel,
			PCTSTR szModule ) 
		{
			tstring sCLSID(_T("CLSID\\"));
			tstring sIndProgID(szProgIDPrefix);
			tstring sProgID;

			// format the prog ID values.
			sIndProgID += tstring(_T(".")) + szClassProgID;
			sProgID = sIndProgID + (TCHAR)'.' + szVersion;

			tstring sGuid(stringFromGuid(rclsid).c_str());
			
			sCLSID += sGuid;

			RegistryKey rkRoot(HKEY_CLASSES_ROOT);
			RegistryKey rk;
			RegistryKey rk2;

			try 
			{
				// create ProgID keys
				rk = rkRoot.createKey(sProgID.c_str());
				rk.setValue(NULL, szDesc);
				rk = rk.createKey(_T("CLSID"));
				rk.setValue(NULL, sGuid.c_str());

				// create VersionIndependentProgID keys.
				rk = rkRoot.createKey(sIndProgID.c_str());
				rk.setValue(NULL, szDesc);
				rk = rk.createKey(_T("CLSID"));
				rk.setValue(NULL, sGuid.c_str());

				// create entries under CLSID.
				rk = rkRoot.createKey(sCLSID.c_str());
				rk.setValue(NULL, szDesc);
				rk2 = rk.createKey(_T("ProgID"));
				rk2.setValue(NULL, sProgID.c_str());
				rk2 = rk.createKey(_T("VersionIndependentProgID"));
				rk2.setValue(NULL, sIndProgID.c_str());
				rk.createKey(_T("NotInsertable"));
				rk2 = rk.createKey(_T("InprocServer32"));
				rk2.setValue(NULL, szModule);
				rk2.setValue(_T("ThreadingModel"), szThreadingModel);
			}
			catch (_win32_error& e)
			{
				// If anything goes wrong, clean-up
				(void)e;
				WXTRACELINE1("error %u creating registry entries", e.error());
				
				unregisterComClass(rclsid, szProgIDPrefix, szVersion, szClassProgID);
				return false;
			}
			
			return true;
		}
			
		static void unregisterComClass( 
			REFCLSID rclsid,		   
			PCTSTR szProgIDPrefix,
			PCTSTR szVersion,			  
			PCTSTR szClassProgID)
		{
			tstring sCLSID(_T("CLSID\\"));
			tstring sIndProgID(szProgIDPrefix);
			tstring sProgID;

			// format the prog ID values.
			sIndProgID += tstring(_T(".")) + szClassProgID;
			sProgID = sIndProgID + (TCHAR)'.' + szVersion;

			tstring sGuid(stringFromGuid(rclsid).c_str());
			
			sCLSID += sGuid;

			RegistryKey rk(HKEY_CLASSES_ROOT);
			RegistryKey rk2;

			#define TRY_CATCH_WIN32_ERROR(x)	try { (x); } catch (_win32_error&) {}
			
			try
			{
				// delete ProgID keys
				rk2 = rk.openKey(sProgID.c_str());
				
				TRY_CATCH_WIN32_ERROR(rk2.deleteKey(_T("CLSID")))
				
				rk.deleteKey(sProgID.c_str());
			}
			catch (_win32_error&) {}

			// delete VersionIndependentProgID keys.
			try
			{
				rk2 = rk.openKey(sIndProgID.c_str());
				
				TRY_CATCH_WIN32_ERROR(rk2.deleteKey(_T("CLSID")));
				
				rk.deleteKey(sIndProgID.c_str());
			}
			catch (_win32_error&) {}

			try
			{
				// delete entries under CLSID.
				rk2 = rk.openKey(sCLSID.c_str());
			
				TRY_CATCH_WIN32_ERROR(rk2.deleteKey(_T("ProgID")));
				TRY_CATCH_WIN32_ERROR(rk2.deleteKey(_T("VersionIndependentProgID")));
				TRY_CATCH_WIN32_ERROR(rk2.deleteKey(_T("NotInsertable")));
				TRY_CATCH_WIN32_ERROR(rk2.deleteKey(_T("InprocServer32")));
				
				rk.deleteKey(sCLSID.c_str());
			}
			catch (_win32_error&){}
		}
	
	public:
		void init(
			HINSTANCE hinst,
			PCTSTR szProgIDPrefix,
			const CO_CLASS_REGISTER* pCoClasses)
		{
			m_hInst = hinst;
			m_progIDPrefix = szProgIDPrefix;
			m_pCoClasses = pCoClasses;
		}

		static HRESULT DllGetClassObject(REFCLSID rclsid, REFIID riid, LPVOID FAR *ppv)				   
		{
			CriticalSection::Auto lock(s_comhelp.m_cs);
			FactoryCollectionIterator classFactory;

			// scan for the right one
			for (const CO_CLASS_REGISTER *pCoClass = ComHelp::s_comhelp.m_pCoClasses; 
				pCoClass->pClsid != NULL; pCoClass++ )
			{
				if (*pCoClass->pClsid == rclsid)
				{
					// If the class factory doesn't exist, add a new one
					if ((classFactory = s_comhelp.m_factory.find((UINT_PTR)pCoClass)) ==
						s_comhelp.m_factory.end())
					{
						ComClassFactory* pClassFactory = new ComClassFactory(pCoClass);
												
						if (pClassFactory == NULL)
            				return E_OUTOFMEMORY;
            				
            			pair<FactoryCollectionIterator, bool> p = 
            				s_comhelp.m_factory.insert(
            					FactoryCollection::value_type(
            						(UINT_PTR)pCoClass, shared_ptr<ComClassFactory>(pClassFactory)));
            						
            			// This means you've got two identical class factories - that's bad
            			_ASSERTE(p.second);
            			
            			classFactory = p.first;

						// Prevent object destruction until shutdown
						reinterpret_cast<IUnknown *>(pClassFactory)->AddRef();
           			}

					return classFactory->second->QueryInterface(riid, ppv);
      			}
			}
				
			return CLASS_E_CLASSNOTAVAILABLE;
		}

		static HRESULT DllRegisterServer()
		{    
			static PCTSTR s_aszModels[] = 
			{
				_T("Apartment"),
				_T("Free"),
				_T("Main"),
				_T("Both"),
				_T("Neutral")	
			};
			HRESULT hr = S_OK;
			TCHAR szModule[_MAX_PATH];	
			
			DllUnregisterServer();
			
			GetModuleFileName(s_comhelp.m_hInst, szModule, ARRAYSIZE(szModule));

			TCHAR szVersion[16];
			bool b = false;
			
			// for each item in the coclass list, register it
			for (const CO_CLASS_REGISTER *pCoClass = s_comhelp.m_pCoClasses; 
				(SUCCEEDED(hr) && (pCoClass->pClsid != NULL)); pCoClass++)
			{
				_itot_s(pCoClass->iVersion, szVersion,16, 10);
				
				// register the class with default values
       			b |= registerComClass( 
       				*pCoClass->pClsid, 
					pCoClass->szCoclassDesc, 
					s_comhelp.m_progIDPrefix.c_str(),
					szVersion, 
					pCoClass->szProgID, 
					s_aszModels[(int)pCoClass->model], 
					szModule);      	        
			}

			return b ? S_OK : E_FAIL;
		}
	
		static HRESULT DllUnregisterServer()
		{    
			TCHAR szVersion[16];
			
			// for each item in the coclass list, unregister it
			for (const CO_CLASS_REGISTER *pCoClass = s_comhelp.m_pCoClasses; 
				pCoClass->pClsid != NULL; pCoClass++ )
			{
				_itot_s(pCoClass->iVersion, szVersion,16, 10);
				
				unregisterComClass(
					*pCoClass->pClsid, 
        			s_comhelp.m_progIDPrefix.c_str(),
					szVersion, 
					pCoClass->szProgID );
			}
		        
			return S_OK;
		}

		static ComHelp s_comhelp;
	private:
    	typedef map<UINT_PTR, shared_ptr<ComClassFactory> > FactoryCollection;
    	typedef FactoryCollection::iterator FactoryCollectionIterator;

		
		
		CriticalSection m_cs;
		HINSTANCE m_hInst;
		tstring m_progIDPrefix;
		const CO_CLASS_REGISTER* m_pCoClasses;
		FactoryCollection m_factory; 
	};

	_declspec(selectany) ComHelp ComHelp::s_comhelp;
	inline static  ComHelp& _stdcall getComHelp()
	{ 
		return ComHelp::s_comhelp; 
	}
};

#pragma warning(pop)
